# Copyright (C) 2022-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).

# --------------------------------------------------------
# CroCo model for downstream tasks
# --------------------------------------------------------

from .croco import CroCoNet

# 从权重中提取参数
def croco_args_from_ckpt(ckpt):
    if 'croco_kwargs' in ckpt: # CroCo v2 released models，有croco_kwargs参数，表示权重是v2模型
        return ckpt['croco_kwargs']
    elif 'args' in ckpt and hasattr(ckpt['args'], 'model'): # pretrained using the official code release，args键中有model属性，表示是使用官方代码的预训练
        s = ckpt['args'].model # eg "CroCoNet(enc_embed_dim=1024, enc_num_heads=16, enc_depth=24)"
        assert s.startswith('CroCoNet(')
        return eval('dict'+s[len('CroCoNet'):]) # transform it into the string of a dictionary and evaluate it，将字符串转换为字典
    else: # CroCo v1 released models
        return dict()

# 下游单目任务编码器
class CroCoDownstreamMonocularEncoder(CroCoNet):
    def __init__(self, head, **kwargs):
        """ Build network for monocular downstream task, only using the encoder.
        It takes an extra argument head, that is called with the features
          and a dictionary img_info containing 'width' and 'height' keys
        The head is setup with the croconet arguments in this init function
        NOTE: It works by *calling super().__init__() but with redefined setters

        """
        super(CroCoDownstreamMonocularEncoder, self).__init__(**kwargs)
        head.setup(self)
        self.head = head

    def _set_mask_generator(self, *args, **kwargs):
        """ No mask generator """
        return

    def _set_mask_token(self, *args, **kwargs):
        """ No mask token """
        self.mask_token = None
        return

    def _set_decoder(self, *args, **kwargs):
        """ No decoder """
        return

    def _set_prediction_head(self, *args, **kwargs):
        """ No 'prediction head' for downstream tasks."""
        return

    def forward(self, img):
        """
        img if of size batch_size x 3 x h x w
        """
        B, C, H, W = img.size()
        img_info = {'height': H, 'width': W}
        need_all_layers = hasattr(self.head, 'return_all_blocks') and self.head.return_all_blocks
        out, _, _ = self._encode_image(img, do_mask=False, return_all_blocks=need_all_layers)
        return self.head(out, img_info)

# 下游双目任务模型，父类是CroCoNet
class CroCoDownstreamBinocular(CroCoNet):
    def __init__(self, head, **kwargs):
        """ Build network for binocular downstream task
        It takes an extra argument head, that is called with the features
          and a dictionary img_info containing 'width' and 'height' keys
        The head is setup with the croconet arguments in this init function
        """
        super(CroCoDownstreamBinocular, self).__init__(**kwargs)
        head.setup(self) # 设置DPT密集预测输出头
        self.head = head # DPT密集预测输出头

    def _set_mask_generator(self, *args, **kwargs):
        """ No mask generator """
        return

    def _set_mask_token(self, *args, **kwargs):
        """ No mask token """
        self.mask_token = None
        return

    def _set_prediction_head(self, *args, **kwargs):
        """ No prediction head for downstream tasks, define your own head """
        return

    def forward(self, img1, img2):
        B, C, H, W = img1.size() # 图像尺寸
        img_info = {'height': H, 'width': W}
        return_all_blocks = hasattr(self.head, 'return_all_blocks') and self.head.return_all_blocks # head中定义，骨干是否需要返回所有块输出(布尔)
        out, pos, _ = self._encode_image(img1, do_mask=False, return_all_blocks=return_all_blocks) # ViT编码器，输入img1，不掩码，返回所有块输出，输出token特征1，补丁位置编码1
        out2, pos2, _ = self._encode_image(img2, do_mask=False, return_all_blocks=False) # 共享权重ViT编码器，输入img2，不掩码，不返回所有块输出，输出token特征2，补丁位置编码2(图像尺寸一样，位置编码一样)
        if return_all_blocks: # 需要返回所有块输出
            decout = self._decoder(out[-1], pos, None, out2, pos2, return_all_blocks=return_all_blocks) # Transformer解码器，输入最后的token特征1，补丁位置编码1，掩码(None)，token特征2，补丁位置编码2，返回所有块输出
            decout = out + decout # 列表连接编码器1(24)和解码器(12)输出(36)
        else:
            decout = self._decoder(out, pos, None, out2, pos2, return_all_blocks=return_all_blocks)
        return self.head(decout, img_info) # DPT密集预测输出头，输入所有块输出(编码器1+解码器，List)，预测图像尺寸
